package com.wry.sharding.algorithm;

import com.wry.sharding.utils.TimeUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.time.LocalDateTime;
import java.util.Collection;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * <h1>  <h1>
 *
 * @author wry
 * @since 2025/1/14 13:40
 */
@Slf4j
public class OrderComplexKeysShardingAlgorithm implements ComplexKeysShardingAlgorithm<Comparable<?>> {


    @Override
    public Collection<String> doSharding(Collection<String> collection, ComplexKeysShardingValue<Comparable<?>> shardingValue) {
        // 1. 获取分片键对应的值
        Map<String, Collection<Comparable<?>>> columnValuesMap = shardingValue.getColumnNameAndShardingValuesMap();

        // pay_time 是 LocalDateTime 类型
        Collection<Comparable<?>> payTimeValues = columnValuesMap.get("pay_time");
        // order_number 是 String 类型
        Collection<Comparable<?>> orderNumberValues = columnValuesMap.get("order_number");

        for (String targetName  : collection) {
            for (Comparable<?> payTimeValue : payTimeValues) {
                LocalDateTime payTime = null;
                if (payTimeValue instanceof String) {
                    payTime = TimeUtil.toLocalDateTime((String) payTimeValue);
                } else if (payTimeValue instanceof LocalDateTime) {
                    payTime = (LocalDateTime) payTimeValue;
                }
                final int month = payTime.getMonth().getValue();
                String tableName = "";
                if (targetName.contains("oms_order_goods_item")) {
                    tableName = "oms_order_goods_item" + payTime.getYear();
                } else {
                    tableName = shardingValue.getLogicTableName() + payTime.getYear();
                }
                if (month < 10) {
                    tableName = tableName + "0" + month;
                } else {
                    tableName = tableName + month;
                }
                if (targetName.contains(tableName)) {
                    log.info("拼接表名 = {}", tableName);
                    return Set.of(tableName);
                }
            }
        }
        return null;
    }
}
